import pygame
from pygame import QUIT, KEYDOWN, K_UP, K_DOWN, K_LEFT, K_RIGHT, K_ESCAPE
import random

ROW_COUNT = 32
COLUMN_COUNT = 24
INIT_POS = (5, 5)
INIT_LENGTH = 8

SCREEN_SIZE = (320, 240)
SCREEN_BG_COLOR = (0, 0, 0)
BALL_COLOR = (255, 255, 255)
BALL_SIZE = 5
SQUARE_SIZE = 5

# const
class MOVE_DIRECTION:
	RIGHT = 1
	DOWN = 2
	LEFT = 3
	UP = 4

class MAP_ELEMENT:
	EMPTY = 0
	SNAKE_BODY = 1
	FOOD = 2

class GamePlay:
	
	def __init__(self):
		self.direction = MOVE_DIRECTION.RIGHT
		self.food_pos = (0, 0)
		#tofix
		self.pos_list = [(INIT_POS[0], INIT_POS[1])] * INIT_LENGTH
		
		self.map = []
		self.initMap()
		
		self.popFood()
		
	def gameReset(self):
		self.__init__()
	
	def initMap(self):
		self.map = []
		
		for i in xrange(COLUMN_COUNT):
			row = []
			for i in xrange(ROW_COUNT):
				row.append(0)
				
			self.map.append(row)
			
	# set all elements to 0
	def clearMap(self):
		for i in xrange(COLUMN_COUNT):
			for j in xrange(ROW_COUNT):
				self.map[i][j] = MAP_ELEMENT.EMPTY
	
	def move(self, dir):
		self.direction = dir
	
	def update(self):
		# update position list
		self.moveSnakeBody()
		
		# put snake on map
		self.updateMap()
		
	def popFood(self):
		x = random.randint(0, ROW_COUNT - 1)
		y = random.randint(0, COLUMN_COUNT - 1)
		
		self.food_pos = (x, y)
		
	def moveSnakeBody(self):
		cur_pos = self.pos_list[0]
		x = cur_pos[0]
		y = cur_pos[1]
		#print cur_pos
		
		# add a ball to next position
		if self.direction == MOVE_DIRECTION.RIGHT:
			x = x + 1 if x + 1 < ROW_COUNT else 0
		elif self.direction == MOVE_DIRECTION.DOWN:
			y = y + 1 if y + 1 < COLUMN_COUNT else 0
		elif self.direction == MOVE_DIRECTION.LEFT:
			x = x - 1 if x - 1 >= 0 else ROW_COUNT - 1
		elif self.direction == MOVE_DIRECTION.UP:
			y = y - 1 if y - 1 >= 0 else COLUMN_COUNT - 1
		
		next_pos = (x, y)
		self.pos_list.insert(0, next_pos)
		
		# do something at next position
		self.triggerAtNextPos(next_pos)
		
	def triggerAtNextPos(self, pos):
		x = pos[0]
		y = pos[1]
		value = self.map[y][x]
		if value == MAP_ELEMENT.EMPTY:
			# remove a ball from rear position
			self.pos_list.pop()
		elif value == MAP_ELEMENT.FOOD:
			self.popFood()
		elif value == MAP_ELEMENT.SNAKE_BODY:
			#todo
			self.gameReset()
		
	def updateMap(self):
		self.clearMap()
		
		for pos in self.pos_list:
			x = pos[0]
			y = pos[1]
			#print (x, y)
			self.map[y][x] = MAP_ELEMENT.SNAKE_BODY
			
		# show food
		self.map[self.food_pos[1]][self.food_pos[0]] = MAP_ELEMENT.FOOD
		
class GameViewer:
	def __init__(self):
		pygame.init()
		
		self.screen = pygame.display.set_mode(SCREEN_SIZE)
		
	def draw(self, map):
		# fill background
		self.screen.fill(SCREEN_BG_COLOR)
		
		# draw map
		for y, row in enumerate(map):
			for x, element in enumerate(row):
				if element == MAP_ELEMENT.SNAKE_BODY:
					self.drawBall((x, y))
				elif element == MAP_ELEMENT.FOOD:
					self.drawSquare((x, y))
		
		pygame.display.flip()
	
	def drawBall(self, map_pos):
		offset = BALL_SIZE
		x = map_pos[0] * BALL_SIZE * 2 + offset
		y = map_pos[1] * BALL_SIZE * 2 + offset
		pygame.draw.circle(self.screen, BALL_COLOR, (x, y), BALL_SIZE)

	def drawSquare(self, map_pos):
		offset = SQUARE_SIZE
		x = map_pos[0] * BALL_SIZE * 2 + offset
		y = map_pos[1] * BALL_SIZE * 2 + offset
		pygame.draw.circle(self.screen, BALL_COLOR, (x, y), SQUARE_SIZE)

		
class GameClient:
	def __init__(self):
		self.view = GameViewer()
		self.game_play = GamePlay()
	
	def run(self):		
		while True:
			self.processInput()
			self.game_play.update()
			self.view.draw(self.game_play.map)
			pygame.time.delay(66)
			
	def processInput(self):
		for event in pygame.event.get():
			if event.type == QUIT:
				exit()
			elif event.type == KEYDOWN:
				if event.key == K_UP:
					self.game_play.move(MOVE_DIRECTION.UP)
				elif event.key == K_DOWN:
					self.game_play.move(MOVE_DIRECTION.DOWN)
				elif event.key == K_LEFT:
					self.game_play.move(MOVE_DIRECTION.LEFT)
				elif event.key == K_RIGHT:
					self.game_play.move(MOVE_DIRECTION.RIGHT)
				elif event.key == K_ESCAPE:
					exit()

if __name__ == "__main__":
	game = GameClient()
	game.run()